#include <windows.h>
#include "CDXBuffer.h"
#include "CReleaseMacro.h"
#include <atlcomcli.h>

CDXBuffer::CDXBuffer(CDXDevice* pDevice) :  CDXResource(pDevice)
{
	m_pBuffer = NULL;
	m_pCPUAccessBuffer = NULL;
	m_szElem = 0;
	m_nElem = 0;
}
CDXBuffer::~CDXBuffer()
{
	ReleaseResources();
}

HRESULT CDXBuffer::ReleaseResources()
{
	SAFE_RELEASE(m_pBuffer);
	SAFE_RELEASE(m_pCPUAccessBuffer);
	return __super::ReleaseResources();
}
ID3D11Resource* CDXBuffer::GetResource() const
{
	return (ID3D11Resource*)m_pBuffer;
}

HRESULT CDXBuffer::Create(D3D11_BUFFER_DESC* pDesc, DWORD szElem, DWORD nElem, D3D11_SUBRESOURCE_DATA* pInitData)
{
	if (m_pDevice == NULL || pDesc == NULL || szElem ==0 || nElem == 0)
		return E_FAIL;	
	ID3D11Device* pDX11Device = m_pDevice->GetD3D11Device();
	HRESULT hr = S_OK;
	ID3D11Buffer * pBuffer = NULL;
	hr = pDX11Device->CreateBuffer( pDesc, pInitData, &pBuffer);
	if( FAILED( hr ) )
		return hr;
	ReleaseResources();
	m_pBuffer = pBuffer;
	m_nElem = nElem;
	m_szElem = szElem;
	return S_OK;
}



HRESULT CDXBuffer::OpenShareHandle(HANDLE hResource)
{
	if (m_pDevice == NULL)
		return E_FAIL;
	if (hResource == 0)
		return E_INVALIDARG;
	if (hResource == m_hShareHandle)
	{
		return S_OK;
	}
	HRESULT hr = S_OK;
	hr = ReleaseResources();
	CComPtr<ID3D11Resource> pResource = NULL;
	ID3D11Device* pDX11Device = m_pDevice->GetD3D11Device();
	hr = pDX11Device->OpenSharedResource(hResource, __uuidof(ID3D11Resource), (void**)(&pResource)); 
	if (FAILED(hr))
		return hr;
	hr = pResource->QueryInterface(__uuidof(ID3D11Buffer), (void**)(&m_pBuffer)); 
	if (FAILED(hr))
		return hr; 
	m_hShareHandle = hResource;
	return S_OK;
}

HRESULT CDXBuffer::Write(BYTE* data, unsigned long szData, unsigned long srcOffset, unsigned long dstOffset
						  ,BOOL bDiscard)
{
	if (m_pDevice == NULL || m_pBuffer == NULL || data == NULL || szData == 0 )
		return E_FAIL;
	ID3D11DeviceContext* pDeviceContext = m_pDevice->GetD3D11DeviceContext();
	HRESULT hr = S_OK;
	ID3D11Buffer* pDstBuf = m_pBuffer;
	D3D11_BUFFER_DESC bufDesc;
	m_pBuffer->GetDesc(&bufDesc);
	if (!(bufDesc.CPUAccessFlags & D3D11_CPU_ACCESS_WRITE))
	{
		if (m_pCPUAccessBuffer == NULL)
		{
			hr = CreateCPUAccessBuffer();
			if (FAILED(hr))
				return hr;
		}
	
		pDstBuf = m_pCPUAccessBuffer;

	}
	if (bufDesc.ByteWidth < szData + dstOffset)
		return E_FAIL;

	D3D11_MAPPED_SUBRESOURCE mappedData;
	D3D11_MAP flag = D3D11_MAP_WRITE;
	if (bDiscard)
	{
		flag = D3D11_MAP_WRITE_DISCARD;
	}

	hr = pDeviceContext->Map(pDstBuf, 0, flag, 0, &mappedData);
	if (FAILED(hr))
		return hr;
	BYTE* pInData = (BYTE*)data;
	BYTE* bufData = (BYTE*)mappedData.pData;
	unsigned long stride = mappedData.RowPitch; 
	for (unsigned long pivot = 0; pivot < szData ; pivot += stride)
	{	
		memcpy_s(&bufData[pivot + dstOffset], stride, &pInData[pivot + srcOffset], min(szData - pivot, stride));
	}
	pDeviceContext->Unmap(pDstBuf, 0);

	if (pDstBuf == m_pCPUAccessBuffer)
	{
		pDeviceContext->CopyResource(m_pBuffer, m_pCPUAccessBuffer);
	}
	return hr;
}

HRESULT CDXBuffer::Read(BYTE* pData, unsigned long szData, unsigned long offset, BOOL bUpdate)
{
	if (pData == NULL || m_pDevice == NULL || m_pBuffer == NULL )
		return E_FAIL;
	ID3D11DeviceContext* pDeviceContext = m_pDevice->GetD3D11DeviceContext();
	HRESULT hr = S_OK;

	ID3D11Buffer* pSrcBuffer = m_pBuffer;
	D3D11_BUFFER_DESC bufDesc;
	m_pBuffer->GetDesc(&bufDesc);

	if (!(bufDesc.CPUAccessFlags & D3D11_CPU_ACCESS_READ))
	{
		if (m_pCPUAccessBuffer == NULL)
		{
			hr = CreateCPUAccessBuffer();
			if (FAILED(hr))
				return hr;
			bUpdate = TRUE;
		}
		if (bUpdate)
		{
			pDeviceContext->CopyResource(m_pCPUAccessBuffer, m_pBuffer);
			if (FAILED(hr))
				return hr;
		}
		pSrcBuffer = m_pCPUAccessBuffer;
	}
	if (szData == 0 || szData + offset > bufDesc.ByteWidth)
	{
		return E_FAIL;
	}

	D3D11_MAPPED_SUBRESOURCE mappedData;
	D3D11_MAP flag = D3D11_MAP_READ;

	hr = pDeviceContext->Map(pSrcBuffer, 0, flag, 0, &mappedData);
	if (FAILED(hr))
		return hr;

	BYTE* bufData = (BYTE*)mappedData.pData;
	unsigned long stride = mappedData.RowPitch; 

	unsigned long pivotDest = 0;
	for (unsigned long pivot = offset; pivot < szData + offset; pivot += stride)
	{	
		UINT numCopied = min(szData + offset - pivot, stride);
		memcpy_s(&pData[pivotDest], stride, &bufData[pivot], numCopied);
		pivotDest += numCopied;
	}

	pDeviceContext->Unmap(pSrcBuffer, 0);
	return hr;
}

ID3D11Buffer* CDXBuffer::GetBuffer()
{
	return m_pBuffer;
}
HRESULT CDXBuffer::GetDesc(D3D11_BUFFER_DESC& bufDesc, UINT& szElem, UINT& nElem)
{
	if (m_pBuffer == NULL)
		return E_FAIL;
	m_pBuffer->GetDesc(&bufDesc);
	szElem = m_szElem;
	nElem = m_nElem;
	return S_OK;
}

HRESULT CDXBuffer::CreateCPUAccessBuffer()
{
	if (m_pDevice == NULL || m_pBuffer == NULL)
		return E_FAIL;
	ID3D11Device* pDX11Device = m_pDevice->GetD3D11Device();
	HRESULT hr = S_OK;
	SAFE_RELEASE(m_pCPUAccessBuffer);
	D3D11_BUFFER_DESC bufDesc;
	m_pBuffer->GetDesc(&bufDesc);

	bufDesc.Usage = D3D11_USAGE_STAGING;
	bufDesc.BindFlags = 0;
	bufDesc.CPUAccessFlags = D3D11_CPU_ACCESS_READ | D3D11_CPU_ACCESS_WRITE;
	bufDesc.MiscFlags = 0;
	hr = pDX11Device->CreateBuffer(&bufDesc, NULL, &m_pCPUAccessBuffer);
	return hr;
}
HRESULT CDXBuffer::CloneTo(CDXBuffer** ppDstBuf)
{
	HRESULT hr = S_OK;
	D3D11_BUFFER_DESC desc;
	UINT nElem =0, szElem = 0;
	hr = this->GetDesc(desc, szElem, nElem);
	if (FAILED(hr))
		return hr;
	BOOL bNeedCreate = TRUE;
	if (*ppDstBuf != NULL)
	{
		D3D11_BUFFER_DESC dstDesc;
		UINT nDstElem =0, szDstElem =0;
		(*ppDstBuf)->GetDesc(dstDesc, szDstElem, nDstElem);
		if (memcmp(&dstDesc, &desc, sizeof(D3D11_BUFFER_DESC)) == 0)
		{
			bNeedCreate = FALSE;
		}
	}
	if (bNeedCreate)
	{
		SAFE_RELEASE(*ppDstBuf);
		*ppDstBuf = new CDXBuffer(m_pDevice);
		hr = (*ppDstBuf)->Create(&desc, szElem, nElem, NULL);
		if (FAILED(hr))
			return hr;
	}
	(*ppDstBuf)->m_szElem = szElem;
	(*ppDstBuf)->m_nElem = nElem;
	hr = (*ppDstBuf)->CopyFrom(this);

	return hr;
}
HRESULT CDXBuffer::GetShaderResourceView(DXGI_FORMAT fmt, ID3D11ShaderResourceView** ppShaderResourceView)
{
	HRESULT hr = S_OK;
	D3D11_SHADER_RESOURCE_VIEW_DESC SRVDesc = {fmt, D3D11_SRV_DIMENSION_BUFFER, {0, m_nElem} };
	hr = GetShaderResourceView(&SRVDesc, ppShaderResourceView);	
	return hr;
}

HRESULT CDXBuffer::GetUnorderedAccessView(DXGI_FORMAT fmt, UINT uavFlag, ID3D11UnorderedAccessView** ppUnorderAccessView)
{
	HRESULT hr = S_OK;
	D3D11_UNORDERED_ACCESS_VIEW_DESC UAVDesc = {fmt, D3D11_UAV_DIMENSION_BUFFER, {0, m_nElem, uavFlag}};
	hr = GetUnorderedAccessView(&UAVDesc, ppUnorderAccessView);
	return hr;
}

HRESULT CDXBuffer::GetRenderTargetView(DXGI_FORMAT fmt, ID3D11RenderTargetView** ppRenderTargetView)
{
	HRESULT hr = S_OK;
	D3D11_RENDER_TARGET_VIEW_DESC RTVDesc = {fmt, D3D11_RTV_DIMENSION_BUFFER, {0, m_szElem}};
	hr = GetRenderTargetView(&RTVDesc, ppRenderTargetView);
	return hr;
}